#!/bin/bash
# ______________________________________________________________________________________________________#
# 3. Docker
# ______________________________________________________________________________________________________#

INTERFACE="eth0"
INTERFACE_VIRTUAL_BRIDGE="virbr1"
COMMAND=0
PAUSE=1
USER=$(whoami)

check_is_root() {
    if [ "$USER" != "root" ]; then
        clear
        echo "Run as root!"
        exit
    fi 
}

print_menu() {
	clear
	echo
	echo "Host Firewall"
	echo "1. Accept ip4v      [80,53,443,67,68, docker0]"
	echo "2. Accept ip4v/ip6v [80,53,443,67,68]"
	echo "3. Block"
	echo "4. Enable log"
	echo "5. Check ruleset"
	echo "6. Flush ruleset"
	echo "7. Add rules for virt-manager"
	echo "0. Exit "
	echo
	return
}

check_ruleset() {
    while [ true ]; do
        sleep 1
        clear
        nft list ruleset
    done
}

get_user_command() {
	echo -n " Please enter an integer -> "
	read COMMAND
}

accept_ip4v() {

    chattr -i /etc/nftables.conf
    #                                       Rules (ip4v)
    # ______________________________________________________________________________________________________#
    nft flush ruleset

    # TABLE FILTER
    nft add table ip filter
    nft add chain ip filter INPUT   { type filter hook input priority   0 \; policy drop \;}
    nft add chain ip filter FORWARD { type filter hook forward priority 0 \; policy drop \;}
    nft add chain ip filter OUTPUT  { type filter hook output priority  0 \; policy drop \;}

    # Chain INPUT
    nft add rule ip filter INPUT iif "lo" counter log prefix "Loopback" accept
    nft add rule ip filter INPUT iif "lo" ip saddr != 127.0.0.0/8 counter log drop
    nft add rule ip filter INPUT iif "docker0" ip saddr 172.17.0.0/16 counter log prefix \"[IN DOCKER]\" accept # DOCKER INPUT ACCEPT
    nft add rule ip filter INPUT iif "$INTERFACE" tcp sport {80, 443, 8080} tcp dport 32768-65535 ct state established counter log prefix \"[IN HTTP/HTTPS]\"       accept
    #nft add rule ip filter INPUT iif "$INTERFACE" tcp sport 22 tcp dport 32768-65535 ct state established counter log prefix \"[IN SSH]\"       accept
    nft add rule ip filter INPUT iif "$INTERFACE" udp sport {67, 68}        udp dport 32768-65535 ct state established counter log prefix \"[IN DHCP]\"             accept
    nft add rule ip filter INPUT iif "$INTERFACE" udp sport 53              udp dport 32768-65535 ct state established counter log prefix \"[IN DNS]\"              accept
    nft add rule ip filter INPUT iif "$INTERFACE" udp sport {45000-60000} udp dport 32768-65535 ct state established counter log prefix \"[IN DISCORD]\"            accept
    nft add rule ip filter INPUT iif "$INTERFACE" icmp type echo-reply                            ct state established counter log prefix \"[IN ICPM ECHO-REPLAY]\" accept
    nft add rule ip filter INPUT meta pkttype { broadcast, multicast }  drop   # DROP broadcast, multicast pkgs
    nft add rule ip filter INPUT ct state invalid                       drop   # DROP invalid pkg
    nft add rule ip filter INPUT icmp type echo-reply                   drop   # DROP ICMP type 0 (Echo Reply)
    nft add rule ip filter INPUT icmp type echo-request                 drop   # DROP ICMP type 8 (Echo Request)
    nft add rule ip filter INPUT counter drop

    # Chain OUTPUT
    nft add rule ip filter OUTPUT oif "lo" counter accept # Accept "Loopback"
    nft add rule ip filter OUTPUT oif "lo" ip daddr != 127.0.0.0/8 counter log drop
    nft add rule ip filter OUTPUT oif "docker0" ip saddr 172.17.0.0/16 counter log prefix \"[OUT DOCKER]\" accept # DOCKER OUTPUT ACCEPT
    nft add rule ip filter OUTPUT oif "$INTERFACE" tcp dport {80, 443, 8080} tcp sport 32768-65535 ct state new, established counter log prefix \"[OUT HTTP/HTTPS]\"        accept
    #nft add rule ip filter OUTPUT oif "$INTERFACE" tcp dport 22 tcp sport 32768-65535 ct state new, established counter log prefix \"[OUT SSH]\"        accept
    nft add rule ip filter OUTPUT oif "$INTERFACE" udp dport {67, 68}        udp sport 32768-65535 ct state new, established counter log prefix \"[OUT DHCP]\"              accept
    nft add rule ip filter OUTPUT oif "$INTERFACE" udp dport 53              udp sport 32768-65535 ct state new, established counter log prefix \"[OUT DNS]\"               accept
    nft add rule ip filter OUTPUT oif "$INTERFACE" udp dport {45000-60000} udp sport 32768-65535 ct state new, established counter log prefix \"[OUT DISCORD]\" accept
    nft add rule ip filter OUTPUT oif "$INTERFACE" icmp type echo-request                                                    counter log prefix \"[OUT ICPM ECHO-REQUEST]\" accept
    nft add rule ip filter OUTPUT meta pkttype { broadcast, multicast }                                                                                                     drop   # DROP broadcast, multicast pkgs
    nft add rule ip filter OUTPUT ct state invalid                                                                                                                          drop   # DROP invalid pkg
    nft add rule ip filter OUTPUT counter drop

    # TABLE NAT
    nft add table ip nat
    nft add chain ip nat prerouting { type nat hook prerouting priority 0 \; }
    nft add chain ip nat postrouting { type nat hook postrouting priority 100 \; }
    nft add rule ip nat postrouting ip saddr 172.17.0.0/16 oif "$INTERFACE" masquerade

    #                                        Rules (ip6v)
    # ______________________________________________________________________________________________________#
    nft add table ip6 filter

    # Default chain's policy
    nft add chain ip6 filter INPUT   { type filter hook input priority   0 \; policy drop \;}
    nft add chain ip6 filter FORWARD { type filter hook forward priority 0 \; policy drop \;}
    nft add chain ip6 filter OUTPUT  { type filter hook output priority  0 \; policy drop \;}

    # Chain INPUT
    nft add rule ip6 filter INPUT iif "lo" counter log prefix "Loopback"  accept
    nft add rule ip6 filter INPUT iif "lo" ip6 daddr != ::1/128 counter log drop
    nft add rule ip6 filter INPUT counter log prefix \"[IN IPV6]\" drop

    # Chain FORWARD
    nft add rule ip6 filter FORWARD counter drop

    # Chain OUTPUT
    nft add rule ip6 filter OUTPUT oif "lo" counter accept # Accept "Loopback"
    nft add rule ip6 filter OUTPUT oif "lo" ip6 daddr != ::1/128 counter log drop
    nft add rule ip6 filter OUTPUT counter drop


}

add_rules_for_virt_manager() {
    nft add table ip filter

    # Default chain's policy
    nft add chain ip filter LIBVIRT_INP   { type filter hook input priority   0 \; policy drop \;}
    nft add chain ip filter LIBVIRT_FORP  { type filter hook forward priority 0 \; policy drop \;}
    nft add chain ip filter LIBVIRT_OUTP  { type filter hook output priority  0 \; policy drop \;}

    # Chain INPUT
    nft add rule ip filter LIBVIRT_INP iif "$INTERFACE_VIRTUAL_BRIDGE" tcp dport 67 counter accept # ALLOW DHCP SERVER FOR VIRT-MANAGER (isolated network)
}

block_inet() {
    #                                       Rules (ip4v, ip6v)
    # ______________________________________________________________________________________________________#
    nft flush ruleset
    nft add table inet filter

    # Default chain's policy
    nft add chain inet filter INPUT   { type filter hook input priority   0 \; policy drop \;}
    nft add chain inet filter FORWARD { type filter hook forward priority 0 \; policy drop \;}
    nft add chain inet filter OUTPUT  { type filter hook output priority  0 \; policy drop \;}

    # Chain INPUT
    nft add rule inet filter INPUT iif "lo" counter log prefix "Loopback"       accept # Accept "Loopback"
    nft add rule inet filter INPUT iif "lo" ip6 daddr != ::1/128    counter log drop
    nft add rule inet filter INPUT iif "lo" ip saddr != 127.0.0.0/8 counter log drop
    nft add rule inet filter INPUT meta pkttype { broadcast, multicast }        drop   # DROP broadcast, multicast pkgs
    nft add rule inet filter INPUT ct state invalid                             drop   # DROP invalid pkg
    nft add rule inet filter INPUT iif "$INTERFACE" icmp type echo-reply        drop   # DROP ICMP type 0 (Echo Reply)
    nft add rule inet filter INPUT iif "$INTERFACE" icmp type echo-request      drop   # DROP ICMP type 8 (Echo Request)
    
    nft add rule inet filter INPUT counter drop

    
    # Chain FORWARD
    nft add rule inet filter FORWARD counter drop

    # Chain OUTPUT
    nft add rule inet filter OUTPUT oif "lo" counter                             accept # Accept "Loopback"
    nft add rule inet filter OUTPUT oif "lo" ip6 daddr != ::1/128    counter log drop
    nft add rule inet filter OUTPUT oif "lo" ip daddr != 127.0.0.0/8 counter log drop
    nft add rule inet filter OUTPUT meta pkttype { broadcast, multicast }        drop   # DROP broadcast, multicast pkgs
    nft add rule inet filter OUTPUT ct state invalid                             drop   # DROP invalid pkg
    nft add rule inet filter OUTPUT counter drop
}

enable_log() {
    echo "net.netfilter.nf_log.2 = nfnetlink_log" >> /etc/sysctl.conf
}

save_rules() {
    echo "#!/usr/sbin/nft -f" > /etc/nftables.conf
    echo "flush ruleset" >> /etc/nftables.conf
    nft list ruleset >> /etc/nftables.conf
    chattr +i /etc/nftables.conf
    # systectl enabale nftables
}

main() {
    check_is_root

	while [ true ]; do
        clear	

		print_menu
		get_user_command
		if [ "$COMMAND" == "0" ]; then 
			clear
			exit
		elif [ "$COMMAND" == "1" ]; then
            accept_ip4v
            save_rules
            nft list ruleset
            sleep $PAUSE
		elif [ "$COMMAND" == "2" ]; then
		    #accept_inet
		    sleep 1
            save_rules
            nft list ruleset
            sleep $PAUSE
        elif [ "$COMMAND" == "3" ]; then
	    block_inet
            save_rules
            nft list ruleset
            sleep $PAUSE
        elif [ "$COMMAND" == "4" ]; then
		    enable_log
            sleep 1
        elif [ "$COMMAND" == "5" ]; then
            check_ruleset
        elif [ "$COMMAND" == "6" ]; then
            chattr -i /etc/nftables.conf
            nft flush ruleset
            save_rules
            sleep 1
        elif [ "$COMMAND" == "7" ]; then
            add_rules_for_virt_manager
            save_rules
            sleep 1
        fi
	done
}

main
